"""
知识库管理服务
负责知识库和文档的管理
"""

from typing import List, Optional, Dict, Any
from sqlalchemy.orm import Session
from sqlalchemy import and_
import os
import uuid
from datetime import datetime

from app.models.knowledge_base import KnowledgeBase, Document
from app.services.base_service import BaseService
from app.schemas.knowledge_base import KnowledgeBaseCreate, KnowledgeBaseUpdate, DocumentCreate
from app.core.config import settings


class KnowledgeBaseService(BaseService[KnowledgeBase, KnowledgeBaseCreate, KnowledgeBaseUpdate]):
    """知识库服务"""
    
    def __init__(self):
        super().__init__(KnowledgeBase)
    
    def get_active_knowledge_bases(self, db: Session) -> List[KnowledgeBase]:
        """获取所有激活的知识库"""
        return db.query(KnowledgeBase).filter(KnowledgeBase.is_active == True).all()
    
    def create_knowledge_base(
        self, 
        db: Session, 
        name: str,
        description: Optional[str] = None,
        embedding_model: str = "text-embedding-ada-002",
        chunk_size: int = 1000,
        chunk_overlap: int = 200
    ) -> KnowledgeBase:
        """创建知识库"""
        # 生成唯一的集合名称，确保只包含Chroma允许的字符
        import re
        # 移除所有非字母数字字符，只保留字母、数字、下划线和连字符
        safe_name = re.sub(r'[^a-zA-Z0-9_-]', '_', name)
        # 确保名称以字母或数字开头和结尾
        safe_name = re.sub(r'^[^a-zA-Z0-9]+', '', safe_name)
        safe_name = re.sub(r'[^a-zA-Z0-9]+$', '', safe_name)
        collection_name = f"kb_{uuid.uuid4().hex[:8]}_{safe_name}"
        
        knowledge_base = KnowledgeBase(
            name=name,
            description=description,
            collection_name=collection_name,
            embedding_model=embedding_model,
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap
        )
        
        db.add(knowledge_base)
        db.commit()
        db.refresh(knowledge_base)
        return knowledge_base
    
    def add_document(
        self, 
        db: Session, 
        knowledge_base_id: int,
        filename: str,
        original_filename: str,
        file_path: str,
        file_size: int,
        file_type: str,
        content: Optional[str] = None
    ) -> Document:
        """添加文档到知识库"""
        document = Document(
            knowledge_base_id=knowledge_base_id,
            filename=filename,
            original_filename=original_filename,
            file_path=file_path,
            file_size=file_size,
            file_type=file_type,
            content=content,
            processing_status="pending"
        )
        
        db.add(document)
        
        # 更新知识库的文档计数
        knowledge_base = db.query(KnowledgeBase).filter(KnowledgeBase.id == knowledge_base_id).first()
        if knowledge_base:
            knowledge_base.document_count += 1
        
        db.commit()
        db.refresh(document)
        return document
    
    def get_knowledge_base_documents(
        self, 
        db: Session, 
        knowledge_base_id: int
    ) -> List[Document]:
        """获取知识库的文档列表"""
        return db.query(Document).filter(Document.knowledge_base_id == knowledge_base_id).all()
    
    def update_document_processing_status(
        self, 
        db: Session, 
        document_id: int, 
        status: str,
        chunk_count: Optional[int] = None,
        error_message: Optional[str] = None
    ) -> bool:
        """更新文档处理状态"""
        document = db.query(Document).filter(Document.id == document_id).first()
        if document:
            document.processing_status = status
            if chunk_count is not None:
                document.chunk_count = chunk_count
            if error_message:
                document.error_message = error_message
            if status == "completed":
                document.is_processed = True
            
            db.commit()
            return True
        return False
    
    def delete_document(self, db: Session, document_id: int) -> bool:
        """删除文档"""
        document = db.query(Document).filter(Document.id == document_id).first()
        if document:
            # 删除物理文件
            if os.path.exists(document.file_path):
                try:
                    os.remove(document.file_path)
                except Exception:
                    pass  # 忽略文件删除错误
            
            # 更新知识库文档计数
            knowledge_base = db.query(KnowledgeBase).filter(
                KnowledgeBase.id == document.knowledge_base_id
            ).first()
            if knowledge_base:
                knowledge_base.document_count -= 1
                knowledge_base.total_chunks -= document.chunk_count or 0
            
            db.delete(document)
            db.commit()
            return True
        return False
    
    def get_processing_documents(self, db: Session) -> List[Document]:
        """获取处理中的文档"""
        return db.query(Document).filter(
            Document.processing_status.in_(["pending", "processing"])
        ).all()
    
    def get_knowledge_base_stats(self, db: Session, knowledge_base_id: int) -> Dict[str, Any]:
        """获取知识库统计信息"""
        knowledge_base = db.query(KnowledgeBase).filter(KnowledgeBase.id == knowledge_base_id).first()
        if not knowledge_base:
            return {}
        
        documents = db.query(Document).filter(Document.knowledge_base_id == knowledge_base_id).all()
        
        return {
            "id": knowledge_base.id,
            "name": knowledge_base.name,
            "description": knowledge_base.description,
            "document_count": knowledge_base.document_count,
            "total_chunks": knowledge_base.total_chunks,
            "embedding_model": knowledge_base.embedding_model,
            "chunk_size": knowledge_base.chunk_size,
            "chunk_overlap": knowledge_base.chunk_overlap,
            "is_active": knowledge_base.is_active,
            "created_at": knowledge_base.created_at,
            "updated_at": knowledge_base.updated_at,
            "documents": [
                {
                    "id": doc.id,
                    "filename": doc.filename,
                    "original_filename": doc.original_filename,
                    "file_size": doc.file_size,
                    "file_type": doc.file_type,
                    "chunk_count": doc.chunk_count,
                    "processing_status": doc.processing_status,
                    "is_processed": doc.is_processed,
                    "created_at": doc.created_at
                }
                for doc in documents
            ]
        }
